{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6ce5fc7d",
   "metadata": {},
   "source": [
    "### Check SV carrier count results "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "15b4fff5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import pandas as pd\n",
    "import pysam\n",
    "from pybedtools import BedTool\n",
    "from io import StringIO\n",
    "from functools import reduce\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7e471646",
   "metadata": {},
   "outputs": [],
   "source": [
    "wkdir = \"/lustre/scratch126/humgen/projects/interval_rna/interval_rna_seq/thomasVDS/misexpression_v3\"\n",
    "wkdir_path = Path(wkdir)\n",
    "\n",
    "chrom=\"chr19\"\n",
    "vcf_path= \"/lustre/scratch126/humgen/projects/interval_rna/interval_rna_seq/thomasVDS/lof_missense/data/sv_vcf/filtered_merged_gs_svp_10728.vcf.gz\"\n",
    "wgs_rna_paired_smpls_path= wkdir_path.joinpath(\"1_rna_seq_qc/wgs_rna_match/paired_wgs_rna_postqc_prioritise_wgs.tsv\")\n",
    "ge_matrix_flat_path = wkdir_path.joinpath(\"2_misexp_qc/misexp_gene_cov_corr/tpm_zscore_4568smpls_8610genes_flat_misexp_corr_qc.csv\")\n",
    "gencode_path = \"/lustre/scratch126/humgen/projects/interval_rna/interval_rna_seq/thomasVDS/misexpression_v3/reference/gencode/gencode.v31.annotation.sorted.gtf.gz\"\n",
    "sv_info_path = \"/lustre/scratch126/humgen/projects/interval_rna/interval_rna_seq/thomasVDS/lof_missense/data/sv_vcf/info_table/final_sites_critical_info_allele.txt\"\n",
    "vep_msc_path = wkdir_path.joinpath(\"4_vrnt_enrich/sv_vep/msc/SV_vep_hg38_msc_parsed.tsv\")\n",
    "vep_all_path = wkdir_path.joinpath(\"4_vrnt_enrich/sv_vep/all/SV_vep_hg38_all_parsed.tsv\")\n",
    "root_dir = wkdir_path.joinpath(\"4_vrnt_enrich/sv_count_carriers/test/windows/200kb_window\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "63088f69",
   "metadata": {},
   "outputs": [],
   "source": [
    "af_cutoff_list = [0, 0.01]\n",
    "af_lower, af_upper = af_cutoff_list\n",
    "window_start = 0\n",
    "window_step_size = 200000 \n",
    "window_max = 1000000 \n",
    "z_cutoff_list = [2, 10, 20, 30, 40]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "503ecf2c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Inputs:\n",
      "- Chromosome: chr19\n",
      "- AF range: 0-0.01\n",
      "- Window start: 0\n",
      "- Window size: 200000\n",
      "- Max window: 1000000\n",
      "- Z-score cutoffs: [2, 10, 20, 30, 40]\n",
      "\n",
      "Gene IDs passing filters: 8610\n",
      "RNA-seq sample IDs passing QC: 4568\n",
      "\n",
      "Loading input VCF ...\n",
      "VCF loaded.\n",
      "Subset VCF to samples with RNA-seq ...\n",
      "Number of samples in VCF with RNA ID and passing QC: 2640\n",
      "Number of RNA IDs passing QC: 2640\n",
      "Number of genes pass QC on chr19: 295\n",
      "Sample IDs in gene expression matrix with SV calls: 2640\n",
      "SV types in SV info file: ['DEL', 'MEI', 'DUP', 'INV']\n"
     ]
    }
   ],
   "source": [
    "print(\"Inputs:\")\n",
    "print(f\"- Chromosome: {chrom}\")\n",
    "print(f\"- AF range: {af_lower}-{af_upper}\")\n",
    "print(f\"- Window start: {window_start}\")\n",
    "print(f\"- Window size: {window_step_size}\")\n",
    "print(f\"- Max window: {window_max}\")\n",
    "print(f\"- Z-score cutoffs: {z_cutoff_list}\")\n",
    "print(\"\")\n",
    "\n",
    "# constants \n",
    "tpm_cutoff = 0.5 \n",
    "\n",
    "# create root directory \n",
    "root_dir_path = Path(root_dir)\n",
    "root_dir_path.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "### Collect samples with VCF calls and RNA sequencing \n",
    "# read in gene expression file \n",
    "ge_matrix_flat_df = pd.read_csv(ge_matrix_flat_path, sep=\",\")\n",
    "gene_id_pass_qc_set = set(ge_matrix_flat_df.gene_id.unique())\n",
    "print(f\"Gene IDs passing filters: {len(gene_id_pass_qc_set)}\")\n",
    "smpl_id_pass_qc_set = set(ge_matrix_flat_df.rna_id.unique())\n",
    "print(f\"RNA-seq sample IDs passing QC: {len(smpl_id_pass_qc_set)}\")\n",
    "print(\"\")\n",
    "# egan ID, RNA ID sample links \n",
    "wgs_rna_paired_smpls_df = pd.read_csv(wgs_rna_paired_smpls_path, sep=\"\\t\")\n",
    "egan_ids_with_rna = wgs_rna_paired_smpls_df[wgs_rna_paired_smpls_df.rna_id.isin(smpl_id_pass_qc_set)].egan_id.tolist()\n",
    "# load VCF and subset to EGAN IDs with RNA \n",
    "print(\"Loading input VCF ...\")\n",
    "vcf_path = vcf_path\n",
    "vcf = pysam.VariantFile(vcf_path, mode = \"r\")\n",
    "print(\"VCF loaded.\")\n",
    "print(\"Subset VCF to samples with RNA-seq ...\")\n",
    "vcf_samples = [sample for sample in vcf.header.samples]\n",
    "vcf_egan_ids_with_rna = set(egan_ids_with_rna).intersection(set(vcf_samples))\n",
    "vcf.subset_samples(vcf_egan_ids_with_rna)\n",
    "vcf_samples_with_rna = [sample for sample in vcf.header.samples]\n",
    "print(f\"Number of samples in VCF with RNA ID and passing QC: {len(vcf_samples_with_rna)}\")\n",
    "# subset egan ID and RNA ID links to samples with SV calls and passing QC \n",
    "wgs_rna_paired_smpls_with_sv_calls_df = wgs_rna_paired_smpls_df[wgs_rna_paired_smpls_df.egan_id.isin(vcf_samples_with_rna)]\n",
    "# write EGAN-RNA ID pairs to file\n",
    "egan_rna_smpls_dir = root_dir_path.joinpath(\"egan_rna_smpls\")\n",
    "Path(egan_rna_smpls_dir).mkdir(parents=True, exist_ok=True)\n",
    "wgs_rna_paired_smpls_with_sv_calls_df.to_csv(egan_rna_smpls_dir.joinpath(\"egan_rna_ids_paired_pass_qc.tsv\"), sep=\"\\t\", index=False)\n",
    "rna_id_pass_qc_sv_calls = wgs_rna_paired_smpls_with_sv_calls_df.rna_id.unique().tolist()\n",
    "print(f\"Number of RNA IDs passing QC: {len(rna_id_pass_qc_sv_calls)}\")\n",
    "\n",
    "### write bed file for genes on chromosome passing QC \n",
    "gene_bed_dir = root_dir_path.joinpath(\"genes_bed\")\n",
    "gene_bed_dir.mkdir(parents=True, exist_ok=True)\n",
    "gene_bed_path = gene_bed_dir.joinpath(f\"{chrom}_genes.bed\")\n",
    "gene_id_pass_qc_on_chrom = []\n",
    "with open(gene_bed_path, \"w\") as f:\n",
    "    for gtf in pysam.TabixFile(gencode_path).fetch(chrom, parser = pysam.asGTF()):\n",
    "        if gtf.feature == \"gene\" and gtf.gene_id.split('.')[0] in gene_id_pass_qc_set:\n",
    "            # check for multiple entries with same name \n",
    "            gene_id_list, chrom_list, start_list, end_list, strand_list = [], [], [], [], []\n",
    "            gene_id_list.append(gtf.gene_id.split('.')[0])\n",
    "            chrom_list.append(gtf.contig)\n",
    "            start_list.append(gtf.start)\n",
    "            end_list.append((gtf.end))\n",
    "            strand_list.append((gtf.strand))\n",
    "            # check or write output\n",
    "            if len(gene_id_list) > 1 or len(chrom_list) > 1 or len(start_list) > 1 or len(end_list) > 1:\n",
    "                print(f\"{gene_id} has multiple entries in gencode file - excluded from output file.\")\n",
    "            elif len(chrom_list) == 0 or len(start_list) == 0 or len(end_list) == 0: \n",
    "                print(f\"{gene_id} has no entries in gencode file - excluded from output file.\")\n",
    "            else: \n",
    "                gene_id, gtf_chrom, start, end, strand = gene_id_list[0], chrom_list[0], start_list[0], end_list[0], strand_list[0]\n",
    "                gene_id_pass_qc_on_chrom.append(gene_id)\n",
    "                chrom_num = gtf_chrom.split(\"chr\")[1]\n",
    "                f.write(f\"{chrom_num}\\t{start}\\t{end}\\t{gene_id}\\t0\\t{strand}\\n\")\n",
    "print(f\"Number of genes pass QC on {chrom}: {len(gene_id_pass_qc_on_chrom)}\")\n",
    "\n",
    "# subset gene expression file to genes on chromosome \n",
    "ge_matrix_flat_chrom_df = ge_matrix_flat_df[ge_matrix_flat_df.gene_id.isin(gene_id_pass_qc_on_chrom)]\n",
    "# subset gene expression file to samples with SV calls \n",
    "ge_matrix_flat_chrom_egan_df = pd.merge(ge_matrix_flat_chrom_df, wgs_rna_paired_smpls_with_sv_calls_df, how=\"inner\", on=\"rna_id\")\n",
    "print(f\"Sample IDs in gene expression matrix with SV calls: {len(ge_matrix_flat_chrom_egan_df.egan_id.unique())}\")\n",
    "# write to file \n",
    "ge_matrix_flat_chrom_egan_dir = root_dir_path.joinpath(\"express_mtx\")\n",
    "Path(ge_matrix_flat_chrom_egan_dir).mkdir(parents=True, exist_ok=True)\n",
    "ge_matrix_flat_chrom_egan_df.to_csv(ge_matrix_flat_chrom_egan_dir.joinpath(f\"{chrom}_ge_mtx_flat.tsv\"), sep=\"\\t\", index=False)\n",
    "# add gene-sample pair to expression dataframe\n",
    "ge_matrix_flat_chrom_egan_df[\"gene_smpl_pair\"] = ge_matrix_flat_chrom_egan_df.gene_id + \",\" + ge_matrix_flat_chrom_egan_df.rna_id\n",
    "\n",
    "### write SVs on chromosome to bed file \n",
    "vrnts_bed_dir = root_dir_path.joinpath(\"vrnts_bed\")\n",
    "vrnts_bed_dir.mkdir(parents=True, exist_ok=True)\n",
    "vrnts_bed_path = vrnts_bed_dir.joinpath(f\"{chrom}_vrnts.bed\")\n",
    "with open(sv_info_path, \"r\") as f_in, open(vrnts_bed_path, \"w\") as f_out:\n",
    "    for line in f_in:\n",
    "        if line.startswith(\"plinkID\"): \n",
    "            continue\n",
    "        else: \n",
    "            vrnt_id, sv_chrom, pos, end = line.split(\"\\t\")[0], line.split(\"\\t\")[2], line.split(\"\\t\")[3], line.split(\"\\t\")[4]\n",
    "            if sv_chrom == chrom:\n",
    "                chrom_num = sv_chrom.split(\"chr\")[1]\n",
    "                f_out.write(f\"{chrom_num}\\t{pos}\\t{end}\\t{vrnt_id}\\n\")\n",
    "\n",
    "### SV information \n",
    "sv_info_df = pd.read_csv(sv_info_path, sep=\"\\t\", dtype={\"plinkID\":str})\n",
    "sv_types_list = sv_info_df.SVTYPE.unique().tolist()\n",
    "print(f\"SV types in SV info file: {sv_types_list}\")\n",
    "sv_info_id_af_df = sv_info_df[[\"plinkID\", \"AF\", \"SVTYPE\"]].rename(columns={\"plinkID\":\"vrnt_id\"})\n",
    "\n",
    "### generate output directories \n",
    "# variant-window intersection directory \n",
    "intersect_bed_dir=root_dir_path.joinpath(\"intersect_bed\")\n",
    "Path(intersect_bed_dir).mkdir(parents=True, exist_ok=True)\n",
    "# genotypes of variants in windows directory \n",
    "intersect_vrnt_gts_dir = root_dir_path.joinpath(\"intersect_vrnt_gts\")\n",
    "Path(intersect_vrnt_gts_dir).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "# load variant and gene bed files \n",
    "vrnts_bed = BedTool(vrnts_bed_path)\n",
    "genes_bed = BedTool(gene_bed_path)\n",
    "# column names for variant window intersection bed file \n",
    "intersect_bed_columns={0:\"chrom_gene\", 1:\"start_gene\", 2:\"end_gene\", 3: \"gene_id\", 4: \"score\", 5: \"strand\", \n",
    "                       6:\"chrom_vrnt\", 7:\"start_vrnt\", 8:\"end_vrnt\", 9:\"vrnt_id\"}\n",
    "\n",
    "# load most severe consequence \n",
    "vep_msc_df = pd.read_csv(vep_msc_path, sep=\"\\t\", dtype={\"vrnt_id\": str})\n",
    "msc_list = vep_msc_df.Consequence.unique().tolist()\n",
    "vep_msc_cnsqn_df = vep_msc_df.rename(columns={\"Uploaded_variation\": \"vrnt_id\"})[[\"vrnt_id\", \"Consequence\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "a13c5948",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Counting variants upstream, 0bp ...\n",
      "Counting variants upstream, 200000bp ...\n",
      "Counting variants upstream, 400000bp ...\n",
      "Counting variants upstream, 600000bp ...\n",
      "Counting variants upstream, 800000bp ...\n",
      "Counting variants upstream, 1000000bp ...\n",
      "Counting variants downstream, 0bp ...\n",
      "Counting variants downstream, 200000bp ...\n",
      "Counting variants downstream, 400000bp ...\n",
      "Counting variants downstream, 600000bp ...\n",
      "Counting variants downstream, 800000bp ...\n",
      "Counting variants downstream, 1000000bp ...\n"
     ]
    }
   ],
   "source": [
    "### assign gene-variant pairs a window  \n",
    "vrnt_gene_pairs_in_windows_df_list = []\n",
    "for direction in [\"upstream\", \"downstream\"]:\n",
    "    window_size = window_start\n",
    "    seen_sv_gene_pairs = set()\n",
    "    while window_size <= window_max:\n",
    "        print(f\"Counting variants {direction}, {window_size}bp ...\")\n",
    "        window_name = f\"{direction}_{window_size}\"\n",
    "        if direction == \"upstream\": \n",
    "            l, r = window_size, 0 \n",
    "        else: \n",
    "            l, r = 0, window_size\n",
    "        vrnts_window_intersect_str = StringIO(str(genes_bed.window(vrnts_bed, \n",
    "                                                                   l=l, \n",
    "                                                                   r=r,\n",
    "                                                                   sw=True\n",
    "                                                                  )))        \n",
    "        intersect_bed_df = pd.read_csv(vrnts_window_intersect_str, sep=\"\\t\", header=None).rename(columns=intersect_bed_columns)\n",
    "        intersect_bed_df = intersect_bed_df[[\"vrnt_id\", \"gene_id\"]].copy()\n",
    "        # variant gene ID column \n",
    "        intersect_bed_df[\"vrnt_gene_id\"] = intersect_bed_df[\"vrnt_id\"].astype(str) + \",\" + intersect_bed_df[\"gene_id\"].astype(str)\n",
    "        # subset to variant gene pairs that have not been seen in previous windows    \n",
    "        sv_gene_pairs_in_window = set(intersect_bed_df.vrnt_gene_id.unique())\n",
    "        new_sv_gene_pairs = sv_gene_pairs_in_window - seen_sv_gene_pairs\n",
    "        new_sv_gene_pairs_df = intersect_bed_df[intersect_bed_df.vrnt_gene_id.isin(new_sv_gene_pairs)].copy()\n",
    "        new_sv_gene_pairs_df[\"window\"] = window_name\n",
    "        vrnt_gene_pairs_in_windows_df_list.append(new_sv_gene_pairs_df)\n",
    "        seen_sv_gene_pairs = seen_sv_gene_pairs.union(sv_gene_pairs_in_window)\n",
    "        window_size += window_step_size\n",
    "\n",
    "vrnt_gene_pairs_in_windows_df = pd.concat(vrnt_gene_pairs_in_windows_df_list)\n",
    "rename_window = {\"upstream_0\": \"gene_body\", \"downstream_0\": \"gene_body\"}\n",
    "vrnt_gene_pairs_in_windows_df.window = vrnt_gene_pairs_in_windows_df.window.replace(rename_window)\n",
    "vrnt_gene_pairs_in_windows_df = vrnt_gene_pairs_in_windows_df.drop_duplicates()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "57d561ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# assign variants MAF bins \n",
    "maf_bins = [0, 0.01, 0.05, 0.1, 0.5]\n",
    "af_bin_labels = [\"0-1\", \"1-5\", \"5-10\", \"10-50\"]\n",
    "sv_info_id_af_df[\"maf_bin\"] = pd.cut(sv_info_id_af_df.AF, bins=maf_bins, labels=af_bin_labels, right=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "b5f5a91f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# add AF and SV type \n",
    "vrnt_gene_pairs_windows_info_df = pd.merge(vrnt_gene_pairs_in_windows_df, \n",
    "                                         sv_info_id_af_df, \n",
    "                                         on=\"vrnt_id\",\n",
    "                                         how=\"inner\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "203a0ed5",
   "metadata": {},
   "outputs": [],
   "source": [
    "### get variant genotypes \n",
    "intersect_vrnt_ids = vrnt_gene_pairs_in_windows_df.vrnt_id.unique()\n",
    "# genotypes for variant IDs in windows\n",
    "vrnt_gt_egan_dict = {}\n",
    "count = 0\n",
    "for vrnt_id in intersect_vrnt_ids:\n",
    "    # get chromosome, start and end of SV \n",
    "    chrom_vrnt, pos, end, = [sv_info_df[sv_info_df.plinkID == vrnt_id][col].item() for col in [\"chr\", \"pos\", \"end\"]]\n",
    "    chrom_vrnt = chrom_vrnt.split(\"chr\")[1]\n",
    "    records = vcf.fetch(str(chrom_vrnt), pos-1, end)\n",
    "    found_vrnt_id = False\n",
    "    for rec in records: \n",
    "        vcf_vrnt_id = str(rec.id)\n",
    "        if vrnt_id == vcf_vrnt_id:\n",
    "            found_vrnt_id = True \n",
    "            gts = [s[\"GT\"] for s in rec.samples.values()]\n",
    "            for i, gt in enumerate(gts): \n",
    "                vrnt_gt_egan_dict[count] = [vrnt_id, vcf.header.samples[i], gt]\n",
    "                count += 1 \n",
    "    if not found_vrnt_id: \n",
    "        raise ValueError(f\"Did not find {vrnt_id} in {vcf_path}\")\n",
    "vrnt_genotypes_df = pd.DataFrame.from_dict(vrnt_gt_egan_dict, orient=\"index\", columns=[\"vrnt_id\", \"egan_id\", \"genotype\"])\n",
    "vrnt_genotypes_df = vrnt_genotypes_df.astype({\"genotype\":str})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "10789146",
   "metadata": {},
   "outputs": [],
   "source": [
    "# add genotypes \n",
    "vrnt_gene_pairs_windows_gts_df = pd.merge(vrnt_gene_pairs_windows_info_df, \n",
    "                                             vrnt_genotypes_df, \n",
    "                                             on=\"vrnt_id\", \n",
    "                                             how=\"inner\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "123883cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# add expression data \n",
    "vrnt_gene_pairs_windows_gts_express_df = pd.merge(vrnt_gene_pairs_windows_gts_df, \n",
    "                                                 ge_matrix_flat_chrom_egan_df, \n",
    "                                                 on=[\"gene_id\", \"egan_id\"],\n",
    "                                                 how=\"inner\"\n",
    "                                                )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c3cd95d4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "vrnt_id             4876\n",
       "gene_id              295\n",
       "vrnt_gene_id       48786\n",
       "window                11\n",
       "AF                   390\n",
       "SVTYPE                 4\n",
       "maf_bin                4\n",
       "egan_id             2640\n",
       "genotype               4\n",
       "rna_id              2640\n",
       "TPM                41142\n",
       "z-score            41466\n",
       "gene_smpl_pair    778800\n",
       "dtype: int64"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# check column numbers \n",
    "#vrnt_gene_pairs_windows_gts_express_df.nunique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "3cdd8e29",
   "metadata": {},
   "outputs": [],
   "source": [
    "# subset to carriers \n",
    "carrier_gts = ['(0, 1)', '(1, 1)']\n",
    "vrnt_gene_carriers_df = vrnt_gene_pairs_windows_gts_express_df[vrnt_gene_pairs_windows_gts_express_df.genotype.isin(carrier_gts)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4172dfe9",
   "metadata": {},
   "source": [
    "### Check gene body window enrichment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "65cafe53",
   "metadata": {},
   "outputs": [],
   "source": [
    "# subset to variants in +/- 200kb and gene body windows \n",
    "gene_body_windows = [\"gene_body\", \"upstream_200000\", \"downstream_200000\"]\n",
    "vrnt_gene_body_window_carriers_df = vrnt_gene_carriers_df[vrnt_gene_carriers_df.window.isin(gene_body_windows)].copy()\n",
    "vrnt_gene_body_window_carriers_df[\"window\"] = \"gene_body_200000\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "5d76e403",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create dataframe with all MAF bin, SV type combinations \n",
    "maf_window_svtype_df = vrnt_gene_body_window_carriers_df[[\"maf_bin\", \"window\", \"SVTYPE\"]].drop_duplicates()\n",
    "maf_window_svtype_df = maf_window_svtype_df[maf_window_svtype_df.maf_bin.notna()].reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "9a32dac1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Z-score threshold: 2\n",
      "\tNumber of misexpressed genes: 150\n",
      "\tNumber of misexpressed gene-sample pairs: 714\n",
      "\tNumber of control gene-sample pairs: 395286\n",
      "Z-score threshold: 10\n",
      "\tNumber of misexpressed genes: 150\n",
      "\tNumber of misexpressed gene-sample pairs: 438\n",
      "\tNumber of control gene-sample pairs: 395562\n",
      "Z-score threshold: 20\n",
      "\tNumber of misexpressed genes: 110\n",
      "\tNumber of misexpressed gene-sample pairs: 162\n",
      "\tNumber of control gene-sample pairs: 290238\n",
      "Z-score threshold: 30\n",
      "\tNumber of misexpressed genes: 54\n",
      "\tNumber of misexpressed gene-sample pairs: 57\n",
      "\tNumber of control gene-sample pairs: 142503\n",
      "Z-score threshold: 40\n",
      "\tNumber of misexpressed genes: 25\n",
      "\tNumber of misexpressed gene-sample pairs: 25\n",
      "\tNumber of control gene-sample pairs: 65975\n"
     ]
    }
   ],
   "source": [
    "count_misexp_cntrl_carriers_gene_body_df_list = []\n",
    "for z_cutoff in z_cutoff_list:\n",
    "    # misexpression events \n",
    "    print(f\"Z-score threshold: {z_cutoff}\")\n",
    "    misexp_df = ge_matrix_flat_chrom_egan_df[(ge_matrix_flat_chrom_egan_df[\"z-score\"] > z_cutoff) & \n",
    "                                             (ge_matrix_flat_chrom_egan_df[\"TPM\"] > tpm_cutoff)]\n",
    "    misexp_genes = misexp_df.gene_id.unique()\n",
    "    print(f\"\\tNumber of misexpressed genes: {len(misexp_genes)}\")\n",
    "    misexp_rna_id = misexp_df.rna_id.unique()\n",
    "    misexp_gene_smpl = misexp_df.gene_smpl_pair.unique()\n",
    "    print(f\"\\tNumber of misexpressed gene-sample pairs: {len(misexp_gene_smpl)}\")\n",
    "    # control events\n",
    "    cntrl_df = ge_matrix_flat_chrom_egan_df[(ge_matrix_flat_chrom_egan_df.gene_id.isin(misexp_genes)) & \n",
    "                                            (~ge_matrix_flat_chrom_egan_df.gene_smpl_pair.isin(misexp_gene_smpl))]\n",
    "    cntrl_gene_smpl = cntrl_df.gene_smpl_pair.unique() \n",
    "    print(f\"\\tNumber of control gene-sample pairs: {len(cntrl_gene_smpl)}\")\n",
    "    # count misexpression SV carriers \n",
    "    misexp_carriers_df = vrnt_gene_body_window_carriers_df[vrnt_gene_body_window_carriers_df.gene_smpl_pair.isin(misexp_gene_smpl)]\n",
    "    count_misexp_carriers_df = pd.DataFrame(misexp_carriers_df.groupby([\"window\", \"maf_bin\"]).gene_smpl_pair.nunique())\n",
    "    count_misexp_carriers_df = count_misexp_carriers_df.reset_index()\n",
    "    count_misexp_carriers_df = count_misexp_carriers_df.rename(columns={\"gene_smpl_pair\": \"all_sv_misexp_check\"})\n",
    "    \n",
    "    # count control SV carriers\n",
    "    cntrl_carriers_df = vrnt_gene_body_window_carriers_df[vrnt_gene_body_window_carriers_df.gene_smpl_pair.isin(cntrl_gene_smpl)]\n",
    "    count_cntrl_carriers_df = pd.DataFrame(cntrl_carriers_df.groupby([\"window\", \"maf_bin\"]).gene_smpl_pair.nunique())\n",
    "    count_cntrl_carriers_df = count_cntrl_carriers_df.reset_index()\n",
    "    count_cntrl_carriers_df = count_cntrl_carriers_df.rename(columns={\"gene_smpl_pair\": \"all_sv_contrl_check\"})\n",
    "    \n",
    "    # add SV type misexpression carriers \n",
    "    count_misexp_carriers_svtype_df = pd.DataFrame(misexp_carriers_df.groupby([\"window\", \"maf_bin\", \"SVTYPE\"]).gene_smpl_pair.nunique())\n",
    "    count_misexp_carriers_svtype_df = count_misexp_carriers_svtype_df.reset_index()\n",
    "    count_misexp_carriers_svtype_df = count_misexp_carriers_svtype_df.pivot(index=[\"window\", \"maf_bin\"], columns=\"SVTYPE\", values=\"gene_smpl_pair\")\n",
    "    count_misexp_carriers_svtype_df = count_misexp_carriers_svtype_df.reset_index()\n",
    "    rename_columns = {sv_type:f\"{sv_type}_misexp_check\" for sv_type in [\"DEL\", \"DUP\", \"INV\", \"MEI\"]}  \n",
    "    count_misexp_carriers_svtype_df = count_misexp_carriers_svtype_df.rename(columns=rename_columns)\n",
    "    \n",
    "    # add SV type control carriers \n",
    "    count_cntrl_carriers_svtype_df = pd.DataFrame(cntrl_carriers_df.groupby([\"window\", \"maf_bin\", \"SVTYPE\"]).gene_smpl_pair.nunique())\n",
    "    count_cntrl_carriers_svtype_df = count_cntrl_carriers_svtype_df.reset_index()\n",
    "    count_cntrl_carriers_svtype_df = count_cntrl_carriers_svtype_df.pivot(index=[\"window\", \"maf_bin\"], columns=\"SVTYPE\", values=\"gene_smpl_pair\")\n",
    "    count_cntrl_carriers_svtype_df = count_cntrl_carriers_svtype_df.reset_index()\n",
    "    rename_columns = {sv_type:f\"{sv_type}_contrl_check\" for sv_type in [\"DEL\", \"DUP\", \"INV\", \"MEI\"]}  \n",
    "    count_cntrl_carriers_svtype_df = count_cntrl_carriers_svtype_df.rename(columns=rename_columns)\n",
    "    \n",
    "    data_frames = [count_misexp_carriers_df, count_cntrl_carriers_df, count_misexp_carriers_svtype_df,\n",
    "                   count_cntrl_carriers_svtype_df]\n",
    "    ### combine all dataframes \n",
    "    df_merged = reduce(lambda  left,right: pd.merge(left,right,on=[\"window\", \"maf_bin\"],\n",
    "                                            how='inner'), data_frames)\n",
    "    \n",
    "    df_merged[\"z_cutoff\"] = z_cutoff\n",
    "    df_merged[\"misexp_genes_check\"] = len(misexp_genes)\n",
    "    df_merged[\"total_misexp_check\"] = len(misexp_gene_smpl)\n",
    "    df_merged[\"total_control_check\"] = len(cntrl_gene_smpl)\n",
    "    count_misexp_cntrl_carriers_gene_body_df_list.append(df_merged)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "2a3d6f33",
   "metadata": {},
   "outputs": [],
   "source": [
    "count_misexp_cntrl_carriers_gene_body_df = pd.concat(count_misexp_cntrl_carriers_gene_body_df_list)\n",
    "count_misexp_cntrl_carriers_gene_body_df.maf_bin = count_misexp_cntrl_carriers_gene_body_df.maf_bin.astype(str)\n",
    "count_misexp_cntrl_carriers_gene_body_df = count_misexp_cntrl_carriers_gene_body_df.fillna(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "2abbb341",
   "metadata": {},
   "outputs": [],
   "source": [
    "# add missing columns \n",
    "\n",
    "sv_carrier_cols = ['all_sv_misexp_check', 'all_sv_contrl_check', 'DEL_misexp_check', 'DUP_misexp_check', \n",
    "                   'MEI_misexp_check', 'INV_misexp_check', 'DEL_contrl_check', 'DUP_contrl_check', \n",
    "                   'INV_contrl_check','MEI_contrl_check']\n",
    "missing_sv_carrier_cols = []\n",
    "for col in sv_carrier_cols: \n",
    "    if col not in count_misexp_cntrl_carriers_gene_body_df.columns: \n",
    "        missing_sv_carrier_cols.append(col)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "825c0cc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "count_misexp_cntrl_carriers_gene_body_all_df = count_misexp_cntrl_carriers_gene_body_df.copy()\n",
    "for col in missing_sv_carrier_cols: \n",
    "    count_misexp_cntrl_carriers_gene_body_all_df[col] = 0 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "590c2d51",
   "metadata": {},
   "outputs": [],
   "source": [
    "# results \n",
    "sv_gene_window_carrier_count_check_path = wkdir_path.joinpath(f\"4_vrnt_enrich/sv_count_carriers/gene_body/200kb_window/carrier_count_gene_msc_reg_af50/{chrom}_carrier_count_gene_msc.tsv\")\n",
    "sv_gene_window_carrier_count_check_df = pd.read_csv(sv_gene_window_carrier_count_check_path, sep=\"\\t\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "0b389ef4",
   "metadata": {},
   "outputs": [],
   "source": [
    "misexp_cntrl_sv_cols = [\"z_cutoff\", \"maf_bin\", \"misexp_genes\", \"all_sv_misexp\", \"total_misexp\", \"all_sv_contrl\", \"total_control\"]\n",
    "misexp_sv_type_cols = [f\"{sv_type}_misexp\" for sv_type in [\"DEL\", \"DUP\", \"INV\", \"MEI\"]]\n",
    "cntrl_sv_type_cols = [f\"{sv_type}_contrl\" for sv_type in [\"DEL\", \"DUP\", \"INV\", \"MEI\"]]\n",
    "cols_to_check = misexp_cntrl_sv_cols + misexp_sv_type_cols + cntrl_sv_type_cols\n",
    "sv_gene_window_carrier_count_trunc_check_df = sv_gene_window_carrier_count_check_df[cols_to_check]\n",
    "sv_gene_window_carrier_count_trunc_z_check_df = sv_gene_window_carrier_count_trunc_check_df[sv_gene_window_carrier_count_trunc_check_df.z_cutoff.isin(z_cutoff_list)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "f2693e59",
   "metadata": {},
   "outputs": [],
   "source": [
    "check_gene_body_carriers_df = pd.merge(count_misexp_cntrl_carriers_gene_body_all_df, \n",
    "                                       sv_gene_window_carrier_count_trunc_z_check_df, \n",
    "                                       on=[\"maf_bin\", \"z_cutoff\"], \n",
    "                                       how=\"inner\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "c5308adb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# checks \n",
    "if not check_gene_body_carriers_df[\"all_sv_misexp\"].equals(check_gene_body_carriers_df[\"all_sv_misexp_check\"]): \n",
    "    raise ValueError(\"Number of misexpression carriers does not match.\")\n",
    "    \n",
    "if not check_gene_body_carriers_df[\"all_sv_contrl_check\"].equals(check_gene_body_carriers_df[\"all_sv_contrl\"]): \n",
    "    raise ValueError(\"Number of control carriers does not match.\")\n",
    "    \n",
    "if not check_gene_body_carriers_df[\"misexp_genes_check\"].equals(check_gene_body_carriers_df[\"misexp_genes\"]): \n",
    "    raise ValueError(\"Number of misexpressed genes does not match.\")\n",
    "    \n",
    "if not check_gene_body_carriers_df[\"total_control_check\"].equals(check_gene_body_carriers_df[\"total_control\"]): \n",
    "    raise ValueError(\"Number of controls does not match.\")\n",
    "    \n",
    "if not check_gene_body_carriers_df[\"total_misexp_check\"].equals(check_gene_body_carriers_df[\"total_misexp\"]): \n",
    "    raise ValueError(\"Number of misexpression events does not match.\")\n",
    "    \n",
    "for sv_type in [\"DEL\", \"DUP\", \"INV\", \"MEI\"]:\n",
    "    for group in [\"misexp\", \"contrl\"]: \n",
    "        if not check_gene_body_carriers_df[f\"{sv_type}_{group}_check\"].astype(int).equals(check_gene_body_carriers_df[f\"{sv_type}_{group}\"].astype(int)): \n",
    "            raise ValueError(f\"Number of {sv_type} {group} events does not match.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a373f419",
   "metadata": {},
   "source": [
    "### Check SV window enrichment "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "cb36fcf4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Z-score threshold: 2\n",
      "\tNumber of misexpressed genes: 150\n",
      "\tNumber of misexpressed gene-sample pairs: 714\n",
      "\tNumber of control gene-sample pairs: 395286\n",
      "Z-score threshold: 10\n",
      "\tNumber of misexpressed genes: 150\n",
      "\tNumber of misexpressed gene-sample pairs: 438\n",
      "\tNumber of control gene-sample pairs: 395562\n",
      "Z-score threshold: 20\n",
      "\tNumber of misexpressed genes: 110\n",
      "\tNumber of misexpressed gene-sample pairs: 162\n",
      "\tNumber of control gene-sample pairs: 290238\n",
      "Z-score threshold: 30\n",
      "\tNumber of misexpressed genes: 54\n",
      "\tNumber of misexpressed gene-sample pairs: 57\n",
      "\tNumber of control gene-sample pairs: 142503\n",
      "Z-score threshold: 40\n",
      "\tNumber of misexpressed genes: 25\n",
      "\tNumber of misexpressed gene-sample pairs: 25\n",
      "\tNumber of control gene-sample pairs: 65975\n"
     ]
    }
   ],
   "source": [
    "count_misexp_cntrl_carriers_df_list = []\n",
    "for z_cutoff in z_cutoff_list:\n",
    "    # misexpression events \n",
    "    print(f\"Z-score threshold: {z_cutoff}\")\n",
    "    misexp_df = ge_matrix_flat_chrom_egan_df[(ge_matrix_flat_chrom_egan_df[\"z-score\"] > z_cutoff) & \n",
    "                                             (ge_matrix_flat_chrom_egan_df[\"TPM\"] > tpm_cutoff)]\n",
    "    misexp_genes = misexp_df.gene_id.unique()\n",
    "    print(f\"\\tNumber of misexpressed genes: {len(misexp_genes)}\")\n",
    "    misexp_rna_id = misexp_df.rna_id.unique()\n",
    "    misexp_gene_smpl = misexp_df.gene_smpl_pair.unique()\n",
    "    print(f\"\\tNumber of misexpressed gene-sample pairs: {len(misexp_gene_smpl)}\")\n",
    "    # control events\n",
    "    cntrl_df = ge_matrix_flat_chrom_egan_df[(ge_matrix_flat_chrom_egan_df.gene_id.isin(misexp_genes)) & \n",
    "                                            (~ge_matrix_flat_chrom_egan_df.gene_smpl_pair.isin(misexp_gene_smpl))]\n",
    "    cntrl_gene_smpl = cntrl_df.gene_smpl_pair.unique() \n",
    "    print(f\"\\tNumber of control gene-sample pairs: {len(cntrl_gene_smpl)}\")\n",
    "    # count misexpression carriers \n",
    "    misexp_carriers_df = vrnt_gene_carriers_df[vrnt_gene_carriers_df.gene_smpl_pair.isin(misexp_gene_smpl)]\n",
    "    count_misexp_carriers_df = pd.DataFrame(misexp_carriers_df.groupby([\"window\", \"maf_bin\"]).gene_smpl_pair.nunique())\n",
    "    count_misexp_carriers_df = count_misexp_carriers_df.reset_index()\n",
    "    count_misexp_carriers_df = count_misexp_carriers_df.rename(columns={\"gene_smpl_pair\": \"all_sv_misexp_check\"})\n",
    "    \n",
    "    # count control carriers\n",
    "    cntrl_carriers_df = vrnt_gene_carriers_df[vrnt_gene_carriers_df.gene_smpl_pair.isin(cntrl_gene_smpl)]\n",
    "    count_cntrl_carriers_df = pd.DataFrame(cntrl_carriers_df.groupby([\"window\", \"maf_bin\"]).gene_smpl_pair.nunique())\n",
    "    count_cntrl_carriers_df = count_cntrl_carriers_df.reset_index()\n",
    "    count_cntrl_carriers_df = count_cntrl_carriers_df.rename(columns={\"gene_smpl_pair\": \"all_sv_contrl_check\"})\n",
    "    \n",
    "    # merge control and misexpression carriers count\n",
    "    count_misexp_cntrl_carriers_df = pd.merge(count_misexp_carriers_df, \n",
    "                                              count_cntrl_carriers_df, \n",
    "                                              on = [\"window\", \"maf_bin\"], \n",
    "                                              how = \"inner\"\n",
    "                                             )\n",
    "    count_misexp_cntrl_carriers_df[\"z_cutoff\"] = z_cutoff\n",
    "    count_misexp_cntrl_carriers_df[\"misexp_genes_check\"] = len(misexp_genes)\n",
    "    count_misexp_cntrl_carriers_df[\"total_misexp_check\"] = len(misexp_gene_smpl)\n",
    "    count_misexp_cntrl_carriers_df[\"total_control_check\"] = len(cntrl_gene_smpl)\n",
    "    count_misexp_cntrl_carriers_df_list.append(count_misexp_cntrl_carriers_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "15ddd9b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "count_misexp_cntrl_carriers_zscore_df = pd.concat(count_misexp_cntrl_carriers_df_list)\n",
    "count_misexp_cntrl_carriers_rare_zscore_df = count_misexp_cntrl_carriers_zscore_df[count_misexp_cntrl_carriers_zscore_df.maf_bin == \"0-1\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "e6a13e2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# results \n",
    "carrier_count_check_path = wkdir_path.joinpath(f\"4_vrnt_enrich/sv_count_carriers/windows/200kb_window/carrier_count/{chrom}_carrier_count.tsv\")\n",
    "carrier_count_check_df = pd.read_csv(carrier_count_check_path, sep=\"\\t\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "07314996",
   "metadata": {},
   "outputs": [],
   "source": [
    "carrier_count_check_df[\"window\"] = carrier_count_check_df.direction + \"_\" + carrier_count_check_df.window_size.astype(str)\n",
    "carrier_count_check_trunc_df = carrier_count_check_df[[\"window\", \"z_cutoff\", \"misexp_genes\", \"all_sv_misexp\", \"total_misexp\", \"all_sv_contrl\", \"total_control\"]]\n",
    "carrier_count_check_trunc_z_df = carrier_count_check_trunc_df[carrier_count_check_trunc_df.z_cutoff.isin(z_cutoff_list)].copy()\n",
    "carrier_count_check_trunc_z_df.window = carrier_count_check_trunc_z_df.window.replace({\"upstream_0\": \"gene_body\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "971649fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "check_carriers_df = pd.merge(carrier_count_check_trunc_z_df, \n",
    "                             count_misexp_cntrl_carriers_rare_zscore_df, \n",
    "                             on=[\"window\", \"z_cutoff\"], \n",
    "                             how=\"inner\")\n",
    "\n",
    "# checks \n",
    "if not check_carriers_df[\"all_sv_misexp\"].equals(check_carriers_df[\"all_sv_misexp_check\"]): \n",
    "    raise ValueError(\"Number of misexpression carriers does not match.\")\n",
    "    \n",
    "if not check_carriers_df[\"all_sv_contrl_check\"].equals(check_carriers_df[\"all_sv_contrl\"]): \n",
    "    raise ValueError(\"Number of control carriers does not match.\")\n",
    "    \n",
    "if not check_carriers_df[\"misexp_genes_check\"].equals(check_carriers_df[\"misexp_genes\"]): \n",
    "    raise ValueError(\"Number of misexpressed genes does not match.\")\n",
    "    \n",
    "if not check_carriers_df[\"total_control_check\"].equals(check_carriers_df[\"total_control\"]): \n",
    "    raise ValueError(\"Number of controls does not match.\")\n",
    "    \n",
    "if not check_carriers_df[\"total_misexp_check\"].equals(check_carriers_df[\"total_misexp\"]): \n",
    "    raise ValueError(\"Number of misexpression events does not match.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5327ac73",
   "metadata": {},
   "source": [
    "### Check VEP consequence enrichment "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "82e7c5a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# annotate variants with VEP consequence in +/- 200kb window \n",
    "# subset to variant-gene carriers in window \n",
    "vrnt_gene_pairs_df = vrnt_gene_body_window_carriers_df[[\"gene_id\", \"vrnt_id\"]].drop_duplicates()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "35d78d0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "vep_msc_ranks = [\n",
    "    \"transcript_ablation\",\n",
    "    \"splice_acceptor_variant\",\n",
    "    \"splice_donor_variant\",\n",
    "    \"stop_gained\",\n",
    "    \"frameshift_variant\",\n",
    "    \"stop_lost\",\n",
    "    \"start_lost\",\n",
    "    \"transcript_amplification\",\n",
    "    \"inframe_insertion\",\n",
    "    \"inframe_deletion\",\n",
    "    \"missense_variant\",\n",
    "    \"protein_altering_variant\",\n",
    "    \"splice_donor_5th_base_variant\", \n",
    "    \"splice_region_variant\",\n",
    "    \"splice_donor_region_variant\", \n",
    "    \"splice_polypyrimidine_tract_variant\"\n",
    "    \"incomplete_terminal_codon_variant\",\n",
    "    \"start_retained_variant\",\n",
    "    \"stop_retained_variant\",\n",
    "    \"synonymous_variant\",\n",
    "    \"coding_sequence_variant\",\n",
    "    \"mature_miRNA_variant\",\n",
    "    \"5_prime_UTR_variant\",\n",
    "    \"3_prime_UTR_variant\",\n",
    "    \"non_coding_transcript_exon_variant\",\n",
    "    \"intron_variant\",\n",
    "    \"NMD_transcript_variant\",\n",
    "    \"non_coding_transcript_variant\",\n",
    "    \"coding_transcript_variant\",\n",
    "    \"upstream_gene_variant\",\n",
    "    \"downstream_gene_variant\",\n",
    "    \"TFBS_ablation\",\n",
    "    \"TFBS_amplification\",\n",
    "    \"TF_binding_site_variant\",\n",
    "    \"regulatory_region_ablation\",\n",
    "    \"regulatory_region_amplification\",\n",
    "    \"feature_elongation\",\n",
    "    \"regulatory_region_variant\",\n",
    "    \"feature_truncation\",\n",
    "    \"intergenic_variant\",\n",
    "    \"sequence_variant\",\n",
    "    \"no_predicted_effect\"]\n",
    "\n",
    "# from all VEP consequences that have no gene annotation\n",
    "vep_msc_not_linked_to_gene = ['TFBS_ablation', \n",
    "                              'TF_binding_site_variant',\n",
    "                              'regulatory_region_variant', \n",
    "                              'TFBS_amplification',\n",
    "                              'intergenic_variant', \n",
    "                              'regulatory_region_ablation',\n",
    "                              'regulatory_region_amplification']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "f5255a5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# annotate variant-gene pairs with VEP consequence \n",
    "vep_all_consq_df = pd.read_csv(vep_all_path, sep=\"\\t\")\n",
    "# collapse to gene-variant consequence pairs \n",
    "vep_all_consq_drop_dups_df = vep_all_consq_df[[\"vrnt_id\", \"gene_id\", \"Consequence\"]].drop_duplicates()\n",
    "# add variant-gene consequences \n",
    "vrnt_gene_pairs_consq_df = pd.merge(vrnt_gene_pairs_df, \n",
    "                                    vep_all_consq_drop_dups_df, \n",
    "                                    on=[\"vrnt_id\", \"gene_id\"], \n",
    "                                    how=\"left\").fillna(\"no_predicted_effect\")\n",
    "# group consequences together for each variant-gene pair\n",
    "vrnt_gene_pairs_consq_df[\"gene_consequence\"] = vrnt_gene_pairs_consq_df.groupby([\"vrnt_id\", \"gene_id\"])['Consequence'].transform(lambda x: ','.join(x))\n",
    "# remove duplicates arising from variant have different consequences across gene transcripts \n",
    "vrnt_gene_pair_consq_collapse_df = vrnt_gene_pairs_consq_df.drop(columns=[\"Consequence\"]).drop_duplicates()\n",
    "if vrnt_gene_pair_consq_collapse_df.shape[0] != vrnt_gene_pairs_df.shape[0]: \n",
    "    raise ValueError(\"Variant gene pair number does not match number of variant gene pairs with consequence.\")\n",
    "# assign each variant-gene pair a unique consequence based on VEP rank  \n",
    "vrnt_gene_pair_msc_consq = []\n",
    "for index, row in vrnt_gene_pair_consq_collapse_df.iterrows(): \n",
    "    gene_consequence = row[\"gene_consequence\"]\n",
    "    gene_consequence_list = gene_consequence.split(\",\")\n",
    "    for consq in vep_msc_ranks: \n",
    "        if consq in gene_consequence_list: \n",
    "            break \n",
    "    vrnt_gene_pair_msc_consq.append(consq)\n",
    "if vrnt_gene_pair_consq_collapse_df.shape[0] != len(vrnt_gene_pair_msc_consq): \n",
    "    raise ValueError(\"Number of MSC per gene does not match number of variant-gene pairs.\")\n",
    "vrnt_gene_pair_consq_collapse_df[\"consequence\"] = vrnt_gene_pair_msc_consq\n",
    "vrnt_gene_pair_consq_collapse_df = vrnt_gene_pair_consq_collapse_df.drop(columns=[\"gene_consequence\"])\n",
    "# split into variants with annotated gene effect and no predicted gene effect\n",
    "vrnt_gene_effect_df = vrnt_gene_pair_consq_collapse_df[vrnt_gene_pair_consq_collapse_df.consequence != \"no_predicted_effect\"]\n",
    "no_predicted_effect_df = vrnt_gene_pair_consq_collapse_df[vrnt_gene_pair_consq_collapse_df.consequence == \"no_predicted_effect\"]\n",
    "# load VEP MSC annotations\n",
    "vep_msc_df = pd.read_csv(vep_msc_path, sep=\"\\t\").rename(columns={\"Uploaded_variation\": \"vrnt_id\", \"Consequence\": \"msc\"})\n",
    "# merge VEP MSC with no predicted effect \n",
    "no_prediced_effect_df = pd.merge(no_predicted_effect_df[[\"vrnt_id\", \"gene_id\"]], \n",
    "                                 vep_msc_df[[\"vrnt_id\", \"msc\"]], \n",
    "                                 on=\"vrnt_id\", \n",
    "                                 how=\"left\")\n",
    "# check for NaNs \n",
    "if no_prediced_effect_df.msc.isnull().values.any(): \n",
    "    raise ValueError(\"Variants missing VEP most severe consequence.\")\n",
    "no_prediced_effect_df[\"consequence\"] = np.where(no_prediced_effect_df.msc.isin(vep_msc_not_linked_to_gene), \n",
    "                                                no_prediced_effect_df.msc, \n",
    "                                                \"no_predicted_effect\")\n",
    "no_prediced_effect_trunc_df = no_prediced_effect_df[[\"vrnt_id\", \"gene_id\", \"consequence\"]]\n",
    "# combine variants with annotated gene effect and variants with no predicted effect with updated regulatory consequence\n",
    "vrnt_gene_pair_consq_msc_added_df = pd.concat([vrnt_gene_effect_df, no_prediced_effect_trunc_df])\n",
    "if vrnt_gene_pair_consq_msc_added_df[[\"vrnt_id\", \"gene_id\"]].drop_duplicates().shape[0] != vrnt_gene_pairs_df.shape[0]: \n",
    "    raise ValueError(\"Number of gene pairs in expression input does not match number of gene pairs with annotated consequence.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "65f53dc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# add variant consequences \n",
    "vrnt_gene_body_window_carriers_vep_df = pd.merge(vrnt_gene_body_window_carriers_df, \n",
    "                                                 vrnt_gene_pair_consq_msc_added_df, \n",
    "                                                 on=[\"vrnt_id\", \"gene_id\"], \n",
    "                                                 how=\"inner\")\n",
    "vrnt_gene_body_window_carriers_vep_df[\"sv_type_consequence\"] = vrnt_gene_body_window_carriers_vep_df[\"SVTYPE\"] + \"_\" + vrnt_gene_body_window_carriers_vep_df.consequence\n",
    "sv_type_consequences = vrnt_gene_body_window_carriers_vep_df.sv_type_consequence.unique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "1c4847be",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Z-score threshold: 2\n",
      "\tNumber of misexpressed genes: 150\n",
      "\tNumber of misexpressed gene-sample pairs: 714\n",
      "\tNumber of control gene-sample pairs: 395286\n",
      "Z-score threshold: 10\n",
      "\tNumber of misexpressed genes: 150\n",
      "\tNumber of misexpressed gene-sample pairs: 438\n",
      "\tNumber of control gene-sample pairs: 395562\n",
      "Z-score threshold: 20\n",
      "\tNumber of misexpressed genes: 110\n",
      "\tNumber of misexpressed gene-sample pairs: 162\n",
      "\tNumber of control gene-sample pairs: 290238\n",
      "Z-score threshold: 30\n",
      "\tNumber of misexpressed genes: 54\n",
      "\tNumber of misexpressed gene-sample pairs: 57\n",
      "\tNumber of control gene-sample pairs: 142503\n",
      "Z-score threshold: 40\n",
      "\tNumber of misexpressed genes: 25\n",
      "\tNumber of misexpressed gene-sample pairs: 25\n",
      "\tNumber of control gene-sample pairs: 65975\n"
     ]
    }
   ],
   "source": [
    "# count carriers \n",
    "count_misexp_cntrl_carriers_gene_body_vep_df_list = []\n",
    "for z_cutoff in z_cutoff_list:\n",
    "    # misexpression events \n",
    "    print(f\"Z-score threshold: {z_cutoff}\")\n",
    "    misexp_df = ge_matrix_flat_chrom_egan_df[(ge_matrix_flat_chrom_egan_df[\"z-score\"] > z_cutoff) & \n",
    "                                             (ge_matrix_flat_chrom_egan_df[\"TPM\"] > tpm_cutoff)]\n",
    "    misexp_genes = misexp_df.gene_id.unique()\n",
    "    print(f\"\\tNumber of misexpressed genes: {len(misexp_genes)}\")\n",
    "    misexp_rna_id = misexp_df.rna_id.unique()\n",
    "    misexp_gene_smpl = misexp_df.gene_smpl_pair.unique()\n",
    "    print(f\"\\tNumber of misexpressed gene-sample pairs: {len(misexp_gene_smpl)}\")\n",
    "    # control events\n",
    "    cntrl_df = ge_matrix_flat_chrom_egan_df[(ge_matrix_flat_chrom_egan_df.gene_id.isin(misexp_genes)) & \n",
    "                                            (~ge_matrix_flat_chrom_egan_df.gene_smpl_pair.isin(misexp_gene_smpl))]\n",
    "    cntrl_gene_smpl = cntrl_df.gene_smpl_pair.unique() \n",
    "    print(f\"\\tNumber of control gene-sample pairs: {len(cntrl_gene_smpl)}\")\n",
    "    \n",
    "    misexp_carriers_df = vrnt_gene_body_window_carriers_vep_df[vrnt_gene_body_window_carriers_vep_df.gene_smpl_pair.isin(misexp_gene_smpl)]\n",
    "    cntrl_carriers_df = vrnt_gene_body_window_carriers_vep_df[vrnt_gene_body_window_carriers_vep_df.gene_smpl_pair.isin(cntrl_gene_smpl)]\n",
    "\n",
    "    # add SV type misexpression carriers \n",
    "    count_misexp_carriers_svtype_df = pd.DataFrame(misexp_carriers_df.groupby([\"window\", \"maf_bin\", \"sv_type_consequence\"]).gene_smpl_pair.nunique())\n",
    "    count_misexp_carriers_svtype_df = count_misexp_carriers_svtype_df.reset_index()\n",
    "    count_misexp_carriers_svtype_df = count_misexp_carriers_svtype_df.pivot(index=[\"window\", \"maf_bin\"], columns=[\"sv_type_consequence\"], values=\"gene_smpl_pair\")\n",
    "    count_misexp_carriers_svtype_df = count_misexp_carriers_svtype_df.reset_index()\n",
    "    rename_columns = {col:f\"{col}_misexp_check\" for col in count_misexp_carriers_svtype_df.columns if col not in [\"window\", \"maf_bin\"]}  \n",
    "    count_misexp_carriers_svtype_df = count_misexp_carriers_svtype_df.rename(columns=rename_columns)\n",
    "    \n",
    "    # add SV type control carriers \n",
    "    count_cntrl_carriers_svtype_df = pd.DataFrame(cntrl_carriers_df.groupby([\"window\", \"maf_bin\",\"sv_type_consequence\"]).gene_smpl_pair.nunique())\n",
    "    count_cntrl_carriers_svtype_df = count_cntrl_carriers_svtype_df.reset_index()\n",
    "    count_cntrl_carriers_svtype_df = count_cntrl_carriers_svtype_df.pivot(index=[\"window\", \"maf_bin\"], columns=[\"sv_type_consequence\"], values=\"gene_smpl_pair\")\n",
    "    count_cntrl_carriers_svtype_df = count_cntrl_carriers_svtype_df.reset_index()\n",
    "    rename_columns = {col:f\"{col}_contrl_check\" for col in count_cntrl_carriers_svtype_df.columns if col not in [\"window\", \"maf_bin\"]}  \n",
    "    count_cntrl_carriers_svtype_df = count_cntrl_carriers_svtype_df.rename(columns=rename_columns)\n",
    "    \n",
    "    data_frames = [count_misexp_carriers_svtype_df, count_cntrl_carriers_svtype_df]\n",
    "    ### combine all dataframes \n",
    "    df_merged = reduce(lambda  left,right: pd.merge(left,right,on=[\"window\", \"maf_bin\"],\n",
    "                                            how='inner'), data_frames)\n",
    "    \n",
    "    df_merged[\"z_cutoff\"] = z_cutoff\n",
    "    df_merged[\"misexp_genes_check\"] = len(misexp_genes)\n",
    "    df_merged[\"total_misexp_check\"] = len(misexp_gene_smpl)\n",
    "    df_merged[\"total_control_check\"] = len(cntrl_gene_smpl)\n",
    "    count_misexp_cntrl_carriers_gene_body_vep_df_list.append(df_merged)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "604e2a9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "count_misexp_cntrl_carriers_gene_body_vep_df = pd.concat(count_misexp_cntrl_carriers_gene_body_vep_df_list)\n",
    "count_misexp_cntrl_carriers_gene_body_vep_df.maf_bin = count_misexp_cntrl_carriers_gene_body_vep_df.maf_bin.astype(str)\n",
    "count_misexp_cntrl_carriers_gene_body_vep_df = count_misexp_cntrl_carriers_gene_body_vep_df.fillna(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "8ebfe2a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# add missing cosequences \n",
    "vep_consq_no_carrier = []\n",
    "for sv_vep in sv_type_consequences:\n",
    "    for group in [\"misexp\", \"contrl\"]:\n",
    "        name = f\"{sv_vep}_{group}_check\"\n",
    "        if name not in count_misexp_cntrl_carriers_gene_body_vep_df.columns: \n",
    "            vep_consq_no_carrier.append(name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "74ac7626",
   "metadata": {},
   "outputs": [],
   "source": [
    "count_misexp_cntrl_carriers_gene_body_vep_all_df = count_misexp_cntrl_carriers_gene_body_vep_df.copy()\n",
    "for vep_consq in vep_consq_no_carrier: \n",
    "    count_misexp_cntrl_carriers_gene_body_vep_all_df[vep_consq] = 0 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "1a338476",
   "metadata": {},
   "outputs": [],
   "source": [
    "check_sv_type_vep_df = pd.merge(sv_gene_window_carrier_count_check_df, \n",
    "                              count_misexp_cntrl_carriers_gene_body_vep_all_df, \n",
    "                              on=[\"maf_bin\", \"z_cutoff\"], \n",
    "                              how=\"inner\"\n",
    "                             )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "2a726310",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sv_vep in sv_type_consequences: \n",
    "    for group in [\"misexp\", \"contrl\"]: \n",
    "        if not check_sv_type_vep_df[f\"{sv_vep}_{group}_check\"].astype(int).equals(check_sv_type_vep_df[f\"{sv_vep}_{group}\"].astype(int)): \n",
    "            raise ValueError(f\"Number of {sv_vep} {group} events does not match.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af8a0030",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
